{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2022 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qFdPvlXBOdUN"
      },
      "source": [
        "# TensorFlow 2 TPUEmbeddingLayer: Quick Start"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/recommenders/examples/tpu_embedding_layer\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/recommenders/blob/main/docs/examples/tpu_embedding_layer.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/recommenders/blob/main/docs/examples/tpu_embedding_layer.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "      \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/recommenders/blob/main/docs/examples/tpu_embedding_layer.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xHxb-dlhMIzW"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This Colab gives a brief introduction into the TPUEmbeddingLayer of TensorFlow 2.\n",
        "\n",
        "The TPUEmbeddingLayer can use the embedding accelerator on the Cloud TPU to speed up embedding lookups when you have many large embedding tables. This is particularly useful when creating recommendation models as these models typically have very large embedding tables.\n",
        "\n",
        "Note: You will need a GCP (Google Compute Engine) account and a GCS (Google Cloud Storage) bucket for this Colab to run.\n",
        "\n",
        "Please follow the [Google Cloud TPU quickstart](https://cloud.google.com/tpu/docs/quickstart) for how to create GCP account and GCS bucket. You have [$300 free credit](https://cloud.google.com/free/) to get started with any GCP product. You can learn more about Cloud TPU at https://cloud.google.com/tpu/docs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUXex9ctTuDB"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Eh-iCRVBm0p"
      },
      "source": [
        "Install Tensorflow 2.0 and Tensorflow-Recommenders"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rEk-ibQkDNtF"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: tensorflow-recommenders in /usr/local/lib/python3.8/dist-packages (0.7.2)\n",
            "Requirement already satisfied: tensorflow\u003e=2.9.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow-recommenders) (2.9.2)\n",
            "Requirement already satisfied: absl-py\u003e=0.1.6 in /usr/local/lib/python3.8/dist-packages (from tensorflow-recommenders) (1.3.0)\n",
            "Requirement already satisfied: numpy\u003e=1.20 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.21.6)\n",
            "Requirement already satisfied: protobuf\u003c3.20,\u003e=3.9.2 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (3.19.6)\n",
            "Requirement already satisfied: astunparse\u003e=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.6.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (57.4.0)\n",
            "Requirement already satisfied: tensorboard\u003c2.10,\u003e=2.9 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (2.9.1)\n",
            "Requirement already satisfied: keras-preprocessing\u003e=1.1.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.1.2)\n",
            "Requirement already satisfied: grpcio\u003c2.0,\u003e=1.24.3 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.51.1)\n",
            "Requirement already satisfied: gast\u003c=0.4.0,\u003e=0.2.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (0.4.0)\n",
            "Requirement already satisfied: tensorflow-io-gcs-filesystem\u003e=0.23.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (0.28.0)\n",
            "Requirement already satisfied: opt-einsum\u003e=2.3.2 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (3.3.0)\n",
            "Requirement already satisfied: keras\u003c2.10.0,\u003e=2.9.0rc0 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (2.9.0)\n",
            "Requirement already satisfied: h5py\u003e=2.9.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (3.1.0)\n",
            "Requirement already satisfied: libclang\u003e=13.0.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (14.0.6)\n",
            "Requirement already satisfied: six\u003e=1.12.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.15.0)\n",
            "Requirement already satisfied: tensorflow-estimator\u003c2.10.0,\u003e=2.9.0rc0 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (2.9.0)\n",
            "Requirement already satisfied: wrapt\u003e=1.11.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.14.1)\n",
            "Requirement already satisfied: flatbuffers\u003c2,\u003e=1.12 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.12)\n",
            "Requirement already satisfied: termcolor\u003e=1.1.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (2.1.1)\n",
            "Requirement already satisfied: typing-extensions\u003e=3.6.6 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (4.4.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (21.3)\n",
            "Requirement already satisfied: google-pasta\u003e=0.1.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (0.2.0)\n",
            "Requirement already satisfied: wheel\u003c1.0,\u003e=0.23.0 in /usr/local/lib/python3.8/dist-packages (from astunparse\u003e=1.6.0-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (0.38.4)\n",
            "Requirement already satisfied: werkzeug\u003e=1.0.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.0.1)\n",
            "Requirement already satisfied: requests\u003c3,\u003e=2.21.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (2.23.0)\n",
            "Requirement already satisfied: google-auth-oauthlib\u003c0.5,\u003e=0.4.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (0.4.6)\n",
            "Requirement already satisfied: google-auth\u003c3,\u003e=1.6.3 in /usr/local/lib/python3.8/dist-packages (from tensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (2.15.0)\n",
            "Requirement already satisfied: tensorboard-plugin-wit\u003e=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.8.1)\n",
            "Requirement already satisfied: tensorboard-data-server\u003c0.7.0,\u003e=0.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (0.6.1)\n",
            "Requirement already satisfied: markdown\u003e=2.6.8 in /usr/local/lib/python3.8/dist-packages (from tensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (3.4.1)\n",
            "Requirement already satisfied: pyasn1-modules\u003e=0.2.1 in /usr/local/lib/python3.8/dist-packages (from google-auth\u003c3,\u003e=1.6.3-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (0.2.8)\n",
            "Requirement already satisfied: cachetools\u003c6.0,\u003e=2.0.0 in /usr/local/lib/python3.8/dist-packages (from google-auth\u003c3,\u003e=1.6.3-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (5.2.0)\n",
            "Requirement already satisfied: rsa\u003c5,\u003e=3.1.4 in /usr/local/lib/python3.8/dist-packages (from google-auth\u003c3,\u003e=1.6.3-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (4.9)\n",
            "Requirement already satisfied: requests-oauthlib\u003e=0.7.0 in /usr/local/lib/python3.8/dist-packages (from google-auth-oauthlib\u003c0.5,\u003e=0.4.1-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.3.1)\n",
            "Requirement already satisfied: importlib-metadata\u003e=4.4 in /usr/local/lib/python3.8/dist-packages (from markdown\u003e=2.6.8-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (4.13.0)\n",
            "Requirement already satisfied: zipp\u003e=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata\u003e=4.4-\u003emarkdown\u003e=2.6.8-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (3.11.0)\n",
            "Requirement already satisfied: pyasn1\u003c0.5.0,\u003e=0.4.6 in /usr/local/lib/python3.8/dist-packages (from pyasn1-modules\u003e=0.2.1-\u003egoogle-auth\u003c3,\u003e=1.6.3-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (0.4.8)\n",
            "Requirement already satisfied: idna\u003c3,\u003e=2.5 in /usr/local/lib/python3.8/dist-packages (from requests\u003c3,\u003e=2.21.0-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,\u003c1.26,\u003e=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests\u003c3,\u003e=2.21.0-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (1.24.3)\n",
            "Requirement already satisfied: certifi\u003e=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests\u003c3,\u003e=2.21.0-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (2022.9.24)\n",
            "Requirement already satisfied: chardet\u003c4,\u003e=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests\u003c3,\u003e=2.21.0-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (3.0.4)\n",
            "Requirement already satisfied: oauthlib\u003e=3.0.0 in /usr/local/lib/python3.8/dist-packages (from requests-oauthlib\u003e=0.7.0-\u003egoogle-auth-oauthlib\u003c0.5,\u003e=0.4.1-\u003etensorboard\u003c2.10,\u003e=2.9-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (3.2.2)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,\u003e=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging-\u003etensorflow\u003e=2.9.0-\u003etensorflow-recommenders) (3.0.9)\n"
          ]
        }
      ],
      "source": [
        "!pip install -U tensorflow-recommenders"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqR2PQG4ZaZ0"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_recommenders as tfrs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l4Av0gMuT9cM"
      },
      "source": [
        "Connect to the TPU node or local TPU and initialize the TPU system.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DnS8adqiUIxl"
      },
      "outputs": [],
      "source": [
        "resolver = tf.distribute.cluster_resolver.TPUClusterResolver('').connect('')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d1d3c1KPoLEC"
      },
      "source": [
        "Create the TPU strategy. Model that needs to run on TPU should be created under TPUStrategy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1KN7A7U_oQvZ"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.TPUStrategy(resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FfpVE4TcpjyI"
      },
      "source": [
        "You can also check the tpu hardware feature in the TPUStrategy object.\n",
        "\n",
        "For example, you can check which version of embedding feature is supported on this TPU. Check the `tf.tpu.experimental.HardwareFeature` for detailed documentation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U2Bzh6DRpQXR"
      },
      "outputs": [],
      "source": [
        "embedding_feature = strategy.extended.tpu_hardware_feature.embedding_feature\n",
        "assert embedding_feature == tf.tpu.experimental.HardwareFeature.EmbeddingFeature.V1, 'Make sure that you have the right TPU Hardware'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kbXNA_dUH2z2"
      },
      "source": [
        "## TPUEmbedding API break down"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yQEg9orXJdy2"
      },
      "source": [
        "### Feature and table configuration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xlheH8_GOWSQ"
      },
      "source": [
        "When creating an instance of this layer, you must specify:\n",
        "\n",
        "1. The complete set of embedding tables,\n",
        "2. The features you expect to lookup in those tables and\n",
        "3. The optimizer(s) you wish to use on the tables.\n",
        "\n",
        "See the documentation of `tf.tpu.experimental.embedding.TableConfig` and `tf.tpu.experimental.embedding.FeatureConfig` for more details on the complete set of options. We will cover the basic usage here."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tkG1xgmVPKLQ"
      },
      "source": [
        "Multiple FeatureConfig objects can use the same TableConfig object, allowing different features to share the same table:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cLqNl8fGJaTh"
      },
      "outputs": [],
      "source": [
        "table_config_one = tf.tpu.experimental.embedding.TableConfig(\n",
        "    vocabulary_size=8, dim=8)\n",
        "table_config_two = tf.tpu.experimental.embedding.TableConfig(\n",
        "    vocabulary_size=16, dim=4)\n",
        "feature_config = {\n",
        "    'feature_one':\n",
        "        tf.tpu.experimental.embedding.FeatureConfig(table=table_config_one),\n",
        "    'feature_two':\n",
        "        tf.tpu.experimental.embedding.FeatureConfig(table=table_config_one),\n",
        "    'feature_three':\n",
        "        tf.tpu.experimental.embedding.FeatureConfig(table=table_config_two)\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L3BJOiMYQPzW"
      },
      "source": [
        "### Optimizers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oVbNemrZQfYu"
      },
      "source": [
        "An optimizer can be globally specified by passing one of the following types of input to the optimizer argument:\n",
        "\n",
        "1. A string, one of 'sgd', 'adagrad' or 'adam', which uses the given optimizer with the default parameters.\n",
        "2. An instance of a Keras optimizer.\n",
        "3. An instance of an optimizer class from the tf.tpu.experimental.embedding module.\n",
        "\n",
        "You may also specify an optimizer at the table level via the optimizer argument of `tf.tpu.experimental.embedding.TableConfig`. This will completely override the global optimizer for this table. For performance reasons it is recommended that you minimize the total number of distinct optimizers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TUjDQlxNPt5S"
      },
      "outputs": [],
      "source": [
        "optimizer=tf.tpu.experimental.embedding.SGD(0.1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S0qyySCgTsFc"
      },
      "source": [
        "### Model Creation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k__tT13pc0rI"
      },
      "source": [
        "Here are two examples of creating a keras model with tpu embedding layer in it.\n",
        "\n",
        "**NOTE**: You can't create two embedding layers on TPU when using the embedding accelerator (EmbeddingFeature.V1 and above)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ySgbe14Lrc_e"
      },
      "source": [
        "For a functional style Keras model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_53v9Bx_Tuj7"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  embedding_inputs = {\n",
        "      'feature_one':\n",
        "          tf.keras.Input(batch_size=1024, shape=(1,), dtype=tf.int32),\n",
        "      'feature_two':\n",
        "          tf.keras.Input(\n",
        "              batch_size=1024, shape=(1,), dtype=tf.int32, ragged=True),\n",
        "      'feature_three':\n",
        "          tf.keras.Input(batch_size=1024, shape=(1,), dtype=tf.int32)\n",
        "  }\n",
        "  # embedding, feature_config and embedding_inputs all have the same nested\n",
        "  # structure.\n",
        "  embedding = tfrs.layers.embedding.TPUEmbedding(\n",
        "      feature_config=feature_config, optimizer=optimizer)(\n",
        "          embedding_inputs)\n",
        "  logits = tf.keras.layers.Dense(1)(\n",
        "      tf.concat(tf.nest.flatten(embedding), axis=1))\n",
        "  model = tf.keras.Model(embedding_inputs, logits)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JFvu93W1rjIO"
      },
      "source": [
        "For a subclass style model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9_ix59m9T5Qd"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:TPU system grpc://10.3.32.50:8470 has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.\n"
          ]
        }
      ],
      "source": [
        "class ModelWithEmbeddings(tf.keras.Model):\n",
        "\n",
        "  def __init__(self):\n",
        "    super(ModelWithEmbeddings, self).__init__()\n",
        "    self.embedding_layer = tfrs.layers.embedding.TPUEmbedding(\n",
        "        feature_config=feature_config, optimizer=optimizer)\n",
        "    self.dense = tf.keras.layers.Dense(1)\n",
        "\n",
        "  def call(self, inputs):\n",
        "    embedding = self.embedding_layer(inputs)\n",
        "    logits = self.dense(tf.concat(tf.nest.flatten(embedding), axis=1))\n",
        "    return logits\n",
        "\n",
        "\n",
        "# Make sure that the tpu is reinitialized when you try to create another mdoel\n",
        "tf.tpu.experimental.initialize_tpu_system(resolver)\n",
        "with strategy.scope():\n",
        "  model = ModelWithEmbeddings()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JWKGIVxrQuyt"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:TPU system grpc://10.3.32.50:8470 has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "\u003ctensorflow.python.tpu.topology.Topology at 0x7f2085f74400\u003e"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tf.tpu.experimental.initialize_tpu_system(resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wPTXbD-cuHk0"
      },
      "source": [
        "## Simple TPUEmbeddingLayer example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tm9gOZv3u5vc"
      },
      "source": [
        "In this tutorial, we build a simple ranking model using the MovieLens 100K dataset with TPUEmbeddingLayer. We can use this model to predict ratings based on `user_id` and `movie_id`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hn0WkEPBvKm0"
      },
      "source": [
        "### Install and import tensorflow datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5Ve0lmwsvP_U"
      },
      "outputs": [],
      "source": [
        "!pip install -q --upgrade tensorflow-datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FSgpqYFRu7aJ"
      },
      "outputs": [],
      "source": [
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zmRwhlC8vhZq"
      },
      "source": [
        "### Read the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-i4UpNd-35Z"
      },
      "source": [
        "In order to make the dataset accessible to the Cloud TPU worker. You need to create a gcs bucket and download the dataset to the bucket. Follow this [instructions](https://cloud.google.com/tpu/docs/storage-buckets) to create your gcs bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyzHwDik_y6H"
      },
      "outputs": [],
      "source": [
        "gcs_bucket = 'gs://YOUR-BUCKET-NAME'  #@param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cU3DyiT_6hVf"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T_sZ245sIrtB"
      },
      "source": [
        "First we fetch the data using tensorflow_dataset. The data that we need is `movie_id`, `user_id` and `user_rating`.\n",
        "\n",
        "Then preprocess the data and convert them to integers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PYL4sPSHvvVM"
      },
      "outputs": [],
      "source": [
        "# Ratings data.\n",
        "ratings = tfds.load(\n",
        "    \"movielens/100k-ratings\", data_dir=gcs_bucket, split=\"train\")\n",
        "\n",
        "# Select the basic features.\n",
        "ratings = ratings.map(\n",
        "    lambda x: {\n",
        "        \"movie_id\": tf.strings.to_number(x[\"movie_id\"]),\n",
        "        \"user_id\": tf.strings.to_number(x[\"user_id\"]),\n",
        "        \"user_rating\": x[\"user_rating\"],\n",
        "    })"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "shOkDa1FgcdL"
      },
      "source": [
        "### Prepare the dataset and model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t9_nYXR4v5C7"
      },
      "source": [
        "Here we define some hyperparameters for the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6eCkL2AUesu6"
      },
      "outputs": [],
      "source": [
        "per_replica_batch_size = 16\n",
        "movie_vocabulary_size = 2048\n",
        "movie_embedding_size = 64\n",
        "user_vocabulary_size = 2048\n",
        "user_embedding_size = 64"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iUuB9QqeLDfU"
      },
      "source": [
        "We'll split the data by putting 80% of the ratings in the train set, and 20% in the test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qNxLkQktK1-H"
      },
      "outputs": [],
      "source": [
        "shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)\n",
        "\n",
        "train = shuffled.take(80_000)\n",
        "test = shuffled.skip(80_000).take(20_000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uJHyGTNEJJoI"
      },
      "source": [
        "Batch the dataset and convert it to a distributed dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n1QON3D1v1Bp"
      },
      "outputs": [],
      "source": [
        "train_dataset = train.batch(\n",
        "    per_replica_batch_size * strategy.num_replicas_in_sync,\n",
        "    drop_remainder=True).cache()\n",
        "test_dataset = test.batch(\n",
        "    per_replica_batch_size * strategy.num_replicas_in_sync,\n",
        "    drop_remainder=True).cache()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HGx0NON0_MQl"
      },
      "outputs": [],
      "source": [
        "distribute_train_dataset = strategy.experimental_distribute_dataset(\n",
        "    train_dataset,\n",
        "    options=tf.distribute.InputOptions(experimental_fetch_to_device=False))\n",
        "distribute_test_dataset = strategy.experimental_distribute_dataset(\n",
        "    test_dataset,\n",
        "    options=tf.distribute.InputOptions(experimental_fetch_to_device=False))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uwaQ6griJW-H"
      },
      "source": [
        "Here we create the optimizer, specify the feature and table config. Then we create the model with embedding layer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8kV9i9q7-mRd"
      },
      "outputs": [],
      "source": [
        "optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.1)\n",
        "\n",
        "user_table = tf.tpu.experimental.embedding.TableConfig(\n",
        "    vocabulary_size=user_vocabulary_size, dim=user_embedding_size)\n",
        "movie_table = tf.tpu.experimental.embedding.TableConfig(\n",
        "    vocabulary_size=movie_vocabulary_size, dim=movie_embedding_size)\n",
        "feature_config = {\n",
        "    \"movie_id\": tf.tpu.experimental.embedding.FeatureConfig(table=movie_table),\n",
        "    \"user_id\": tf.tpu.experimental.embedding.FeatureConfig(table=user_table)\n",
        "}\n",
        "\n",
        "\n",
        "# Define a ranking model with embedding layer.\n",
        "class EmbeddingModel(tfrs.models.Model):\n",
        "\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "\n",
        "    self.embedding_layer = tfrs.layers.embedding.TPUEmbedding(\n",
        "        feature_config=feature_config, optimizer=optimizer)\n",
        "    self.ratings = tf.keras.Sequential([\n",
        "        # Learn multiple dense layers.\n",
        "        tf.keras.layers.Dense(256, activation=\"relu\"),\n",
        "        tf.keras.layers.Dense(64, activation=\"relu\"),\n",
        "        # Make rating predictions in the final layer.\n",
        "        tf.keras.layers.Dense(1)\n",
        "    ])\n",
        "    self.task: tf.keras.layers.Layer = tfrs.tasks.Ranking(\n",
        "        loss=tf.keras.losses.MeanSquaredError(\n",
        "            reduction=tf.keras.losses.Reduction.NONE),\n",
        "        metrics=[tf.keras.metrics.RootMeanSquaredError()])\n",
        "\n",
        "  def compute_loss(self, features, training=False):\n",
        "    embedding = self.embedding_layer({\n",
        "        \"user_id\": features[\"user_id\"],\n",
        "        \"movie_id\": features[\"movie_id\"]\n",
        "    })\n",
        "    rating_predictions = self.ratings(\n",
        "        tf.concat([embedding[\"user_id\"], embedding[\"movie_id\"]], axis=1))\n",
        "\n",
        "    return tf.reduce_sum(\n",
        "        self.task(\n",
        "            labels=features[\"user_rating\"], predictions=rating_predictions)) * (\n",
        "                1 / (per_replica_batch_size * strategy.num_replicas_in_sync))\n",
        "\n",
        "  def call(self, features, serving_config=None):\n",
        "    embedding = self.embedding_layer(\n",
        "        {\n",
        "            \"user_id\": features[\"user_id\"],\n",
        "            \"movie_id\": features[\"movie_id\"]\n",
        "        },\n",
        "        serving_config=serving_config)\n",
        "    return self.ratings(\n",
        "        tf.concat([embedding[\"user_id\"], embedding[\"movie_id\"]], axis=1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ojYAbLd-L0LU"
      },
      "source": [
        "Make sure that you initialize the model under TPUStrategy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1rePm7itLvvN"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  model = EmbeddingModel()\n",
        "  model.compile(optimizer=optimizer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DqbT3zejg1WT"
      },
      "source": [
        "### Train and eval the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eaULpEhEnxtH"
      },
      "outputs": [],
      "source": [
        "import os"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wgrZ7MYOMJFb"
      },
      "source": [
        "Train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uQzSbMAc-XMa"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1/10\n",
            "10/10 [==============================] - 7s 32ms/step - root_mean_squared_error: 2.7897 - loss: 0.0564 - regularization_loss: 0.0000e+00 - total_loss: 0.0564\n",
            "Epoch 2/10\n",
            "10/10 [==============================] - 0s 26ms/step - root_mean_squared_error: 1.1963 - loss: 0.0088 - regularization_loss: 0.0000e+00 - total_loss: 0.0088\n",
            "Epoch 3/10\n",
            "10/10 [==============================] - 0s 25ms/step - root_mean_squared_error: 1.1261 - loss: 0.0089 - regularization_loss: 0.0000e+00 - total_loss: 0.0089\n",
            "Epoch 4/10\n",
            "10/10 [==============================] - 0s 35ms/step - root_mean_squared_error: 1.1403 - loss: 0.0094 - regularization_loss: 0.0000e+00 - total_loss: 0.0094\n",
            "Epoch 5/10\n",
            "10/10 [==============================] - 0s 40ms/step - root_mean_squared_error: 1.1269 - loss: 0.0103 - regularization_loss: 0.0000e+00 - total_loss: 0.0103\n",
            "Epoch 6/10\n",
            "10/10 [==============================] - 0s 36ms/step - root_mean_squared_error: 1.1162 - loss: 0.0100 - regularization_loss: 0.0000e+00 - total_loss: 0.0100\n",
            "Epoch 7/10\n",
            "10/10 [==============================] - 0s 36ms/step - root_mean_squared_error: 1.1365 - loss: 0.0097 - regularization_loss: 0.0000e+00 - total_loss: 0.0097\n",
            "Epoch 8/10\n",
            "10/10 [==============================] - 0s 47ms/step - root_mean_squared_error: 1.1171 - loss: 0.0110 - regularization_loss: 0.0000e+00 - total_loss: 0.0110\n",
            "Epoch 9/10\n",
            "10/10 [==============================] - 0s 48ms/step - root_mean_squared_error: 1.1037 - loss: 0.0100 - regularization_loss: 0.0000e+00 - total_loss: 0.0100\n",
            "Epoch 10/10\n",
            "10/10 [==============================] - 0s 51ms/step - root_mean_squared_error: 1.0953 - loss: 0.0092 - regularization_loss: 0.0000e+00 - total_loss: 0.0092\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "\u003ckeras.callbacks.History at 0x7f2084d7ddf0\u003e"
            ]
          },
          "execution_count": 24,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model.fit(distribute_train_dataset, steps_per_epoch=10, epochs=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NoJIMm7wMFcz"
      },
      "source": [
        "Evaluate the model on test dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RkPVa0v8PcX1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "10/10 [==============================] - 4s 27ms/step - root_mean_squared_error: 1.1339 - loss: 0.0090 - regularization_loss: 0.0000e+00 - total_loss: 0.0090\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "[1.1338995695114136, 0.009662957862019539, 0, 0.009662957862019539]"
            ]
          },
          "execution_count": 25,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model.evaluate(distribute_test_dataset, steps=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9iKIz9Q-g7zq"
      },
      "source": [
        "### Save and restore the checkpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9kgfgJyovBxB"
      },
      "source": [
        "You can use a gcs bucket to store your checkpoint.\n",
        "\n",
        "Make sure that you give the tpu worker access to the bucket by following the [instructions](https://cloud.google.com/tpu/docs/storage-buckets)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zh4X0cB7vAeB"
      },
      "outputs": [],
      "source": [
        "model_dir = os.path.join(gcs_bucket, \"saved_model\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EllhvC2t_3_Q"
      },
      "source": [
        "Create the checkpoint for the TPU model and save the model to the bucket.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cy3D16YKs39a"
      },
      "outputs": [],
      "source": [
        "checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)\n",
        "saved_tpu_model_path = checkpoint.save(os.path.join(model_dir, \"ckpt\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q_5oyoKTAKlQ"
      },
      "source": [
        "You can list the variables that get stored in that path."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MO3-k0vEs430"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[('_CHECKPOINTABLE_OBJECT_GRAPH', []),\n",
              " ('model/embedding_layer/_tpu_embedding/.ATTRIBUTES/TPUEmbedding_saveable',\n",
              "  []),\n",
              " ('model/embedding_layer/_tpu_embedding/table_0/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [2048, 64]),\n",
              " ('model/embedding_layer/_tpu_embedding/table_0/.OPTIMIZER_SLOT/optimizer/accumulator/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [2048, 64]),\n",
              " ('model/embedding_layer/_tpu_embedding/table_1/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [2048, 64]),\n",
              " ('model/embedding_layer/_tpu_embedding/table_1/.OPTIMIZER_SLOT/optimizer/accumulator/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [2048, 64]),\n",
              " ('model/ratings/layer_with_weights-0/bias/.ATTRIBUTES/VARIABLE_VALUE', [256]),\n",
              " ('model/ratings/layer_with_weights-0/bias/.OPTIMIZER_SLOT/optimizer/accumulator/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [256]),\n",
              " ('model/ratings/layer_with_weights-0/kernel/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [128, 256]),\n",
              " ('model/ratings/layer_with_weights-0/kernel/.OPTIMIZER_SLOT/optimizer/accumulator/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [128, 256]),\n",
              " ('model/ratings/layer_with_weights-1/bias/.ATTRIBUTES/VARIABLE_VALUE', [64]),\n",
              " ('model/ratings/layer_with_weights-1/bias/.OPTIMIZER_SLOT/optimizer/accumulator/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [64]),\n",
              " ('model/ratings/layer_with_weights-1/kernel/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [256, 64]),\n",
              " ('model/ratings/layer_with_weights-1/kernel/.OPTIMIZER_SLOT/optimizer/accumulator/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [256, 64]),\n",
              " ('model/ratings/layer_with_weights-2/bias/.ATTRIBUTES/VARIABLE_VALUE', [1]),\n",
              " ('model/ratings/layer_with_weights-2/bias/.OPTIMIZER_SLOT/optimizer/accumulator/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [1]),\n",
              " ('model/ratings/layer_with_weights-2/kernel/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [64, 1]),\n",
              " ('model/ratings/layer_with_weights-2/kernel/.OPTIMIZER_SLOT/optimizer/accumulator/.ATTRIBUTES/VARIABLE_VALUE',\n",
              "  [64, 1]),\n",
              " ('model/task/_ranking_metrics/0/count/.ATTRIBUTES/VARIABLE_VALUE', []),\n",
              " ('model/task/_ranking_metrics/0/total/.ATTRIBUTES/VARIABLE_VALUE', []),\n",
              " ('optimizer/decay/.ATTRIBUTES/VARIABLE_VALUE', []),\n",
              " ('optimizer/iter/.ATTRIBUTES/VARIABLE_VALUE', []),\n",
              " ('optimizer/learning_rate/.ATTRIBUTES/VARIABLE_VALUE', []),\n",
              " ('save_counter/.ATTRIBUTES/VARIABLE_VALUE', [])]"
            ]
          },
          "execution_count": 28,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tf.train.list_variables(saved_tpu_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jy37ojRqASzf"
      },
      "source": [
        "You can restore the checkpoint later. This is a common practice to checkpoint your model for every epoch and restore that afterwards."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lpuEFC4ftSEZ"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  checkpoint.restore(saved_tpu_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dlAuXHvtAf93"
      },
      "source": [
        "Addtionally you can create a cpu model and restore the weights that gets trained on TPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SuKjhxuzt8pE"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "\u003ctensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f20830fe5b0\u003e"
            ]
          },
          "execution_count": 30,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "cpu_model = EmbeddingModel()\n",
        "\n",
        "# Create the cpu checkpoint and restore the tpu checkpoint.\n",
        "cpu_checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=cpu_model)\n",
        "cpu_checkpoint.restore(saved_tpu_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N4VmtDxTCdat"
      },
      "source": [
        "You can also restore embedding weights partially."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jwBuRXIKCmqd"
      },
      "outputs": [],
      "source": [
        "embedding_checkpoint = tf.train.Checkpoint(embedding=model.embedding_layer)\n",
        "saved_embedding_path = embedding_checkpoint.save(\n",
        "    os.path.join(model_dir, 'tpu-embedding'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D_8La-5XDJ3Y"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "\u003ctensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f20831bbeb0\u003e"
            ]
          },
          "execution_count": 32,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Restore the embedding parameters on cpu model.\n",
        "cpu_embedding_checkpoint = tf.train.Checkpoint(\n",
        "    embeddign=cpu_model.embedding_layer)\n",
        "cpu_embedding_checkpoint.restore(saved_embedding_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TaQlUL2RQHHZ"
      },
      "outputs": [],
      "source": [
        "# Save the embedding parameters on cpu model and restore it to the tpu model.\n",
        "saved_cpu_embedding_path = embedding_checkpoint.save(\n",
        "    os.path.join(model_dir, 'cpu-embedding'))\n",
        "with strategy.scope():\n",
        "  embedding_checkpoint.restore(saved_cpu_embedding_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7HFzfTP-CV29"
      },
      "source": [
        "### Serving"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EuAnDhn8F4kr"
      },
      "source": [
        "Finally, You can use the exported cpu model to do serving. Serving is\n",
        "accomplished through the `tf.saved_model` API"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AO2JYrJfH6eB"
      },
      "source": [
        "**Note**: It is important that the input_signature is specified here so that the\n",
        "exported graph has the correct shapes and types. Moreover the function should be\n",
        "a new, untraced function, to allow tf.saved_model.save to make a fresh trace of\n",
        "the function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bbhyRI-M4VSm"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:Skipping full serialization of Keras layer \u003ctensorflow_recommenders.tasks.ranking.Ranking object at 0x7f20831ead00\u003e, because it is not built.\n"
          ]
        }
      ],
      "source": [
        "@tf.function\n",
        "def serve_tensors(features):\n",
        "  return cpu_model(features)\n",
        "\n",
        "\n",
        "signatures = {\n",
        "    'serving':\n",
        "        serve_tensors.get_concrete_function(\n",
        "            features={\n",
        "                'movie_id':\n",
        "                    tf.TensorSpec(shape=(1,), dtype=tf.int32, name='movie_id'),\n",
        "                'user_id':\n",
        "                    tf.TensorSpec(shape=(1,), dtype=tf.int32, name='user_id'),\n",
        "            }),\n",
        "}\n",
        "tf.saved_model.save(\n",
        "    cpu_model,\n",
        "    export_dir=os.path.join(model_dir, 'exported_model'),\n",
        "    signatures=signatures)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nU9Ko8NaIBTw"
      },
      "source": [
        "The exported model can now be loaded (in Python or C) and used for serving"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zkzuJ2i43CTL"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:Detecting that an object or model or tf.train.Checkpoint is being deleted with unrestored values. See the following logs for the specific values in question. To silence these warnings, use `status.expect_partial()`. See https://www.tensorflow.org/api_docs/python/tf/train/Checkpoint#restorefor details about the status object returned by the restore function.\n",
            "WARNING:tensorflow:An attribute in the restored object could not be found in the checkpoint. Object: (root).embedding_layer._tpu_embedding, attribute: ['TPUEmbedding_saveable']\n"
          ]
        }
      ],
      "source": [
        "imported = tf.saved_model.load(os.path.join(model_dir, 'exported_model'))\n",
        "predict_fn = imported.signatures['serving']\n",
        "\n",
        "# Dummy serving data.\n",
        "input_batch = {\n",
        "    'movie_id': tf.constant(np.array([100]), dtype=tf.int32),\n",
        "    'user_id': tf.constant(np.array([30]), dtype=tf.int32)\n",
        "}\n",
        "# The prediction it generates.\n",
        "prediction = predict_fn(**input_batch)['output_0']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e0a9ZsIQW6I3"
      },
      "source": [
        "Additionally, you can pass the serving config to do serving.\n",
        "\n",
        "Note that: You can use a subset of the trained embedding tables to do serving by using a serving config."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BdUfeXXaW3UU"
      },
      "outputs": [],
      "source": [
        "serving_config = {\n",
        "    'movie_id': tf.tpu.experimental.embedding.FeatureConfig(table=movie_table),\n",
        "    'user_id': tf.tpu.experimental.embedding.FeatureConfig(table=user_table)\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XvR1onYQYvu5"
      },
      "outputs": [],
      "source": [
        "prediction = cpu_model(input_batch, serving_config=serving_config)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "tpu_embedding_layer.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
